using Unity.Burst; // 引入 Burst 编译器，用于高性能 job 系统编译
using Unity.Collections; // 引入 Unity 的 Native 容器，比如 NativeArray
using Unity.Jobs; // 引入 Unity Job System，用于多线程任务调度
using Unity.Mathematics; // 引入 Unity 的数学库，提供高效的数学类型与函数

using static Unity.Mathematics.math; // 静态导入 math 命名空间，可以直接使用如 sin, cos, float4 等而不加前缀

public static class Shapes
{
	/// <summary>
	/// 一个包含 4 个点（顶点）信息的结构体，每个点有位置和法线，使用 float4x3 表示 4 个点的 3D 向量
	/// </summary>
	public struct Point4
	{
		public float4x3 positions, normals; // 每个都是 4 行 3 列，代表 4 个点的位置/法线向量
	}

	/// <summary>
	/// 形状的通用接口，定义如何根据索引获取一组 4 个点的位置和法线信息
	/// </summary>
	public interface IShape
	{
		Point4 GetPoint4(int i, float resolution, float invResolution);
	}

	/// <summary>
	/// 平面形状的实现
	/// </summary>
	public struct Plane : IShape
	{
		/// <summary>
		/// 根据索引 i 返回平面上 4 个点的位置和法线
		/// 这些点位于一个平铺的网格上，法线始终朝上（Y 轴正方向）
		/// </summary>
		public Point4 GetPoint4(int i, float resolution, float invResolution)
		{
			// 将索引 i 转换为 4 个 UV 坐标，代表该索引对应的 4 个点在 UV 空间中的位置
			float4x2 uv = IndexTo4UV(i, resolution, invResolution);
			return new Point4
			{
				// 位置：将 UV 映射到 [-0.5, 0.5] 范围内的 XZ 平面，Y 坐标为 0
				positions = float4x3(uv.c0 - 0.5f, 0f, uv.c1 - 0.5f),
				// 法线：所有点的法线均为 Y 轴正方向
				normals = float4x3(0f, 1f, 0f)
			};
		}
	}

	/// <summary>
	/// 球体形状的实现
	/// 使用一种基于 UV 的近似球体生成方式（类似八面体映射或平面映射变形为球体）
	/// </summary>
	public struct Sphere : IShape
	{
		public Point4 GetPoint4(int i, float resolution, float invResolution)
		{
			// 获取当前索引对应的 4 个点的 UV 坐标
			float4x2 uv = IndexTo4UV(i, resolution, invResolution);

			Point4 p;

			// 先将 UV 坐标偏移到 [-0.5, 0.5] 范围，作为球体表面的初始投影
			p.positions.c0 = uv.c0 - 0.5f; // X 方向
			p.positions.c1 = uv.c1 - 0.5f; // Y 方向
										   // Z 方向通过公式计算，使得 X 和 Y 的绝对值之和不超过 0.5，从而形成一个“方块被压扁成球”的效果
			p.positions.c2 = 0.5f - math.abs(p.positions.c0) - math.abs(p.positions.c1);

			// 计算超出球面的部分（负 Z 的部分），并用 offset 记录需要“挤压”的距离
			float4 offset = math.max(-p.positions.c2, 0f);

			// 对于 Z 值为负的点，调整 X 和 Y 的位置，使其向内收缩，保证所有点都在球面内
			p.positions.c0 += math.select(-offset, offset, p.positions.c0 < 0f);
			p.positions.c1 += math.select(-offset, offset, p.positions.c1 < 0f);

			// 对位置进行归一化（缩放至单位球面上）
			// 计算每个点到原点的距离的倒数，用于缩放向量到单位长度
			float4 scale = 0.5f * rsqrt(
				p.positions.c0 * p.positions.c0 +
				p.positions.c1 * p.positions.c1 +
				p.positions.c2 * p.positions.c2
			);
			p.positions.c0 *= scale;
			p.positions.c1 *= scale;
			p.positions.c2 *= scale;

			// 法线等于位置向量（对于单位球体来说，法线即位置方向）
			p.normals = p.positions;

			return p;
		}
	}

	/// <summary>
	/// 环面（圆环 / 甜甜圈）形状的实现
	/// 使用参数方程来生成环面上的点
	/// </summary>
	public struct Torus : IShape
	{
		public Point4 GetPoint4(int i, float resolution, float invResolution)
		{
			// 获取当前索引对应的 4 个点的 UV 坐标
			float4x2 uv = IndexTo4UV(i, resolution, invResolution);

			// 定义环面的两个半径：主半径 r1，次半径 r2
			float r1 = 0.375f;
			float r2 = 0.125f;

			// 计算环面主环上的半径（随 UV 的角度变化）
			float4 s = r1 + r2 * cos(2f * PI * uv.c1); // r1 是主半径，r2*cos(...) 是次半径在 Y 轴上的投影

			Point4 p;

			// 根据环面的参数方程计算 3D 坐标
			p.positions.c0 = s * sin(2f * PI * uv.c0); // X：主半径与角度相关的圆周运动
			p.positions.c1 = r2 * sin(2f * PI * uv.c1); // Y：次半径的垂直分量
			p.positions.c2 = s * cos(2f * PI * uv.c0); // Z：主半径与角度相关的圆周运动

			// 法线初始化为位置向量（暂未归一化，后面会处理）
			p.normals = p.positions;

			// 修正法线，减去环面中心的影响，得到真正的表面法线
			p.normals.c0 -= r1 * sin(2f * PI * uv.c0); // X 方向减去主环的影响
			p.normals.c2 -= r1 * cos(2f * PI * uv.c0); // Z 方向减去主环的影响

			return p;
		}
	}

	/// <summary>
	/// Burst 编译的 Job，用于并行计算多个形状的点位置和法线，并应用变换矩阵
	/// </summary>
	[BurstCompile(FloatPrecision.Standard, FloatMode.Fast, CompileSynchronously = true)]
	public struct Job<S> : IJobFor where S : struct, IShape
	{
		// 输出：每个索引对应的 4 个点的位置（3x4 矩阵，每列是一个 float3 点）
		[WriteOnly] NativeArray<float3x4> positions, normals;

		// 分辨率相关
		public float resolution, invResolution;

		// 变换矩阵：用于对点和法线进行空间变换
		public float3x4 positionTRS, normalTRS;

		/// <summary>
		/// 将 4x3 的向量组转换为 3x4 矩阵（每列一个点）
		/// </summary>
		float4x3 TransformVectors(float3x4 trs, float4x3 p, float w = 1f)
			=> float4x3(
				trs.c0.x * p.c0 + trs.c1.x * p.c1 + trs.c2.x * p.c2 + trs.c3.x * w,
				trs.c0.y * p.c0 + trs.c1.y * p.c1 + trs.c2.y * p.c2 + trs.c3.y * w,
				trs.c0.z * p.c0 + trs.c1.z * p.c1 + trs.c2.z * p.c2 + trs.c3.z * w
			);

		/// <summary>
		/// 每个线程执行的逻辑：计算单个索引 i 对应的 4 个点，并应用变换
		/// </summary>
		public void Execute(int i)
		{
			// 获取该索引对应的 4 个点信息
			Point4 p = default(S).GetPoint4(i, resolution, invResolution);

			// 变换位置向量，并转置为 3x4 矩阵存储
			positions[i] = transpose(TransformVectors(positionTRS, p.positions));

			// 变换法线向量（w=0，避免位移影响），并归一化后存储
			float3x4 n = transpose(TransformVectors(normalTRS, p.normals, 0f));
			normals[i] = float3x4(
				normalize(n.c0), normalize(n.c1), normalize(n.c2), normalize(n.c3)
			);
		}

		/// <summary>
		/// 静态方法：调度 Job，用于并行处理所有点
		/// </summary>
		public static JobHandle ScheduleParallel(
			NativeArray<float3x4> positions, NativeArray<float3x4> normals,
			int resolution, float4x4 trs, JobHandle dependency
		)
		{
			// 计算逆转置矩阵，用于正确变换法线向量
			float4x4 tim = transpose(inverse(trs));
			return new Job<S>
			{
				positions = positions,
				normals = normals,
				resolution = resolution,
				invResolution = 1f / resolution,
				positionTRS = float3x4(trs.c0.xyz, trs.c1.xyz, trs.c2.xyz, trs.c3.xyz),
				normalTRS = float3x4(tim.c0.xyz, tim.c1.xyz, tim.c2.xyz, tim.c3.xyz)
			}.ScheduleParallel(positions.Length, resolution, dependency);
		}
	}

	/// <summary>
	/// 委托：用于统一不同形状的 Job 调度方法
	/// </summary>
	public delegate JobHandle ScheduleDelegate(
		NativeArray<float3x4> positions, NativeArray<float3x4> normals,
		int resolution, float4x4 trs, JobHandle dependency
	);

	/// <summary>
	/// 工具方法：将一个整数索引 i 转换为 4 个 UV 坐标，用于表示该索引对应的 4 个点在 UV 空间中的位置
	/// </summary>
	public static float4x2 IndexTo4UV(int i, float resolution, float invResolution)
	{
		float4x2 uv;
		// 将索引 i 拆分为 4 个子索引：0,1,2,3
		float4 i4 = 4f * i + float4(0f, 1f, 2f, 3f);
		// 计算这些子索引在 UV 空间中的行号
		uv.c1 = floor(invResolution * i4 + 0.00001f);
		// 计算列号，加上 0.5 是为了居中
		uv.c0 = invResolution * (i4 - resolution * uv.c1 + 0.5f);
		uv.c1 = invResolution * (uv.c1 + 0.5f);
		return uv;
	}
}